# Packet结构体定义
import asyncio
import websockets
import json
import threading
import time
import base64
import io
import numpy as np
from typing import Dict, List, Union
import queue
import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent.parent / "debugger"))
from save_file import save_img, save_depth, save_pcd, save_json


class Packet:
    def __init__(self, packet_type: str, data: Union[Dict, List]):
        self.type = packet_type
        self.data = data
    def to_dict(self):
        return {"type": self.type, "data": self.data}
    @classmethod
    def from_dict(cls, packet_dict: Dict):
        return cls(packet_dict["type"], packet_dict["data"])

class WebSocketComm:
    async def request_get_obs_async(self) -> Packet:
        print("[WebSocketComm] 调用 request_get_obs_async")
        packet = Packet("GET_OBS", {})
        response = await self.send_request_and_wait_response_async(packet)
        return response

    async def request_take_action_async(self, action: np.ndarray) -> Packet:
        print(f"[WebSocketComm] 调用 request_take_action_async, action={action}")
        action_data = {"action": list(action)}
        packet = Packet("TAKE_ACTION", action_data)
        response = await self.send_request_and_wait_response_async(packet)
        return response
    def __init__(self, host="127.0.0.1", port=8765):
        self.host = host
        self.port = port
        self.url = f"ws://{host}:{port}"
        self._ws = None
        self._loop = None
        self._ws_ready = threading.Event()
        self._recv_queue = queue.Queue()
        self._server_thread = threading.Thread(target=self._run_server, daemon=True)
        self._server_thread.start()
        self.save_idx = 0

    def _ensure_loop(self):
        if self._loop is None:
            self._loop = asyncio.new_event_loop()
            asyncio.set_event_loop(self._loop)
        return self._loop

    async def _handler(self, connection):
        self._ws = connection
        self._ws_ready.set()
        print(f"[WebSocketComm] 客户端连接成功: {connection.remote_address}")
        try:
            while True:
                msg = await connection.recv()
                packet_dict = json.loads(msg)
                packet = Packet.from_dict(packet_dict)
                self._recv_queue.put(packet)
        except Exception as e:
            self._ws = None
            self._ws_ready.clear()
            print(e.__traceback__)
            exit(-1)

    def _run_server(self):
        print(f"[WebSocketComm] WebSocket服务端启动: ws://{self.host}:{self.port}")
        async def server_main():
            async with websockets.serve(self._handler, self.host, self.port, max_size=None):
                await asyncio.Future()  # 永远阻塞，直到线程退出
        asyncio.run(server_main())

    def _wait_for_connection(self):
        self._ws_ready.wait()

    async def send_request_and_wait_response_async(self, packet: Packet) -> Packet:
        print(f"[WebSocketComm] 调用 send_request_and_wait_response_async, packet={packet.to_dict()}")
        wait_count = 0
        while self._ws is None:
            if wait_count % 100 == 0:
                print("[WebSocketComm] waiting for client connection...")
            wait_count += 1
            await asyncio.sleep(0.01)
        data = packet.to_dict()
        serialized = json.dumps(data, default=str)
        await self._ws.send(serialized)
        # 等待队列返回响应
        while True:
            try:
                response_packet = self._recv_queue.get(timeout=1)
                break
            except queue.Empty:
                print("[WebSocketComm] waiting for response from client...")
        return response_packet

    def recv_packet(self, block=True, timeout=None) -> Packet:
        print(f"[WebSocketComm] 调用 recv_packet, block={block}, timeout={timeout}")
        start = time.time()
        while True:
            try:
                packet = self._recv_queue.get(block=False)
                return packet
            except queue.Empty:
                pass
            if not block:
                return None
            if timeout is not None and (time.time() - start) > timeout:
                return None
            time.sleep(0.01)

    def send_packet(self, packet: Packet):
        print(f"[WebSocketComm] 调用 send_packet, packet={packet.to_dict()}")
        self._wait_for_connection()
        loop = self._ensure_loop()
        async def _send():
            data = packet.to_dict()
            serialized = json.dumps(data, default=str)
            await self._ws.send(serialized)
        return loop.run_until_complete(_send())

    def close(self):
        if self._ws:
            loop = self._ensure_loop()
            loop.run_until_complete(self._ws.close())
            self._ws = None
        if self._loop:
            self._loop.close()
            self._loop = None

    def dumps_npz_b64(self, arrays):
        buffer = io.BytesIO()
        np.savez_compressed(buffer, **arrays)
        return base64.b64encode(buffer.getvalue()).decode('utf-8')

    def loads_npz_b64(self, b64_string):
        raw_data = base64.b64decode(b64_string)
        buffer = io.BytesIO(raw_data)
        with np.load(buffer, allow_pickle=False) as data:
            return {key: data[key] for key in data.files}
        return base64.b64encode(buffer.getvalue()).decode('utf-8')

    def get_obs(self, idx=None, no_log=False):
        if self._ws is None:
            raise ValueError("websocket_comm 未初始化")
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = None
        if loop and loop.is_running():
            future = asyncio.run_coroutine_threadsafe(self.request_get_obs_async(), loop)
            response = future.result()
        else:
            response = asyncio.run(self.request_get_obs_async())
        response_dict = response.to_dict()
        if not response_dict["data"]["success"]:
            print(f"[WebSocketComm] 获取观察失败: {response_dict['data']['error']}")
            return None
        observation_data = response_dict["data"]["obs"]
        observation = {}
        for key, value in observation_data.items():
            if isinstance(value, (list, np.ndarray)):
                observation[key] = np.array(value)
            else:
                observation[key] = value
        print(f"[WebSocketComm] 成功获取观察: {list(observation.keys())}")
        if not no_log:
            save_idx = idx if idx is not None else self.save_idx
            save_dir = Path(__file__).parent.parent / "piper-imgs"
            if "image" in observation:
                print(f"  - image shape: {observation['image'].shape}")
                save_img(str(save_dir / f"test_image-{save_idx}.png"), observation['image'])
            if "depth" in observation:
                print(f"  - depth shape: {observation['depth'].shape}")
                save_depth(str(save_dir / f"test_depth-{save_idx}.png"), observation['depth'])
            if "point_cloud" in observation:
                print(f"  - point_cloud shape: {observation['point_cloud'].shape}")
                save_pcd(str(save_dir / f"test_point_cloud-{save_idx}.pcd"), observation['point_cloud'])
            if "joint_state" in observation:
                print(f"  - joint_state shape: {observation['joint_state'].shape}")
                save_json(str(save_dir / f"test_joint_state-{save_idx}.json"), observation['joint_state'])
            if "ee_state" in observation:
                print(f"  - ee_state shape: {observation['ee_state'].shape}")
                save_json(str(save_dir / f"test_ee_state-{save_idx}.json"), observation['ee_state'])
            if idx is None:
                self.save_idx += 1
        return observation

    def take_action(self, action):
        if self._ws is None:
            raise ValueError("websocket_comm 未初始化")
        if isinstance(action, np.ndarray):
            action_list = action.tolist()
        else:
            action_list = list(action)
        async def send_action_via_websocket(action_list):
            response = await self.request_take_action_async(np.array(action_list))
            return response.to_dict()
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = None
        if loop and loop.is_running():
            future = asyncio.run_coroutine_threadsafe(send_action_via_websocket(action_list), loop)
            response_dict = future.result()
        else:
            response_dict = asyncio.run(send_action_via_websocket(action_list))
        if not response_dict["data"]["success"]:
            print(f"[WebSocketComm] 执行动作失败: {response_dict['data']['error']}")
            return False
        print(f"[WebSocketComm] 成功执行动作: {len(action_list)} 维动作向量")
        return True
